import re
import html
from datetime import datetime

TITLE_ALIASES = {
    "vuln": [r"漏洞识别", r"漏洞发现", r"Vulnerability Identification", r"Vulnerability Findings"],
    "impact": [r"影响评估", r"影响分析", r"Impact Assessment", r"Impact Analysis"],
    "reco": [r"安全建议", r"修复建议", r"Mitigation", r"Recommendations", r"Security Recommendations"],
}

NORM_TITLES = {
    "vuln": "## 漏洞识别",
    "impact": "## 影响评估",
    "reco": "## 安全建议",
}

def _norm_newlines(s: str) -> str:
    return s.replace("\r\n", "\n").replace("\r", "\n")

def _fix_heading_spaces(s: str) -> str:
    """
    把 '##影响评估' 统一修复为 '## 影响评估'
    同时确保标题后换行（标题行末尾没有内容）
    """
    s = re.sub(r"(?m)^\s*##\s*(\S)", lambda m: "## " + m.group(1), s)
    # 把同一行里紧跟正文的情况拆开（确保标题行只有标题）
    s = re.sub(r"(?m)^(## [^\n#]+)\s+(?!#)(.+)$", r"\1\n\2", s)
    return s

def _trim_md_backticks(s: str) -> str:
    # 把 ` withdraw ` → `withdraw`
    return re.sub(r"`\s*([^`]+?)\s*`", r"`\1`", s)

def _ensure_code_fences(s: str) -> str:
    # 统一为 ```solidity 围栏
    def repl(m):
        body = m.group(1)
        return "```solidity\n" + body.strip("\n") + "\n```"
    return re.sub(r"```(?:sol|solidity)?\s*\n([\s\S]*?)```", repl, s)

def _clean_list_indentation(s: str) -> str:
    lines = s.split("\n")
    out = []
    for line in lines:
        line = line.replace("•", "-")
        # 统一二级缩进
        line = re.sub(r"^\s{3,}-\s+", "  - ", line)
        # 统一一级条目
        line = re.sub(r"^\s*-\s+", "- ", line)
        out.append(line)
    return "\n".join(out)

def _find_first_section(text: str, aliases) -> str | None:
    """
    找到匹配别名的第一个 '## 标题' 区块（直到下一个 '## ' 或文本结束）
    返回纯内容（不含标题），若找不到返回 None
    """
    # 转义每个别名中的特殊字符，然后组合
    escaped_aliases = [re.escape(alias) for alias in aliases]
    title_pat = "|".join(escaped_aliases)
    # 允许 ## 标题前后出现空白；区块到下一节或文本末尾
    # 使用 re.IGNORECASE | re.MULTILINE | re.DOTALL 标志
    pattern = rf"^\s*##\s*(?:{title_pat})\s*\n(.*?)(?=^\s*##\s+|\Z)"
    m = re.search(pattern, text, flags=re.IGNORECASE | re.MULTILINE | re.DOTALL)
    if not m:
        return None
    return m.group(1).strip()

def extract_security_md(analysis_text: str) -> str:
    """
    抽取并规范化【漏洞识别 / 影响评估 / 安全建议】三节 Markdown。
    - 同名章节只保留第一次出现（自动去重）
    - 修复 '##影响评估' 无空格 & 标题行后缺少换行问题
    - 标题统一标准化为中文（## 漏洞识别 / ## 影响评估 / ## 安全建议）
    - 代码块围栏统一为 ```solidity
    - 修复行内反引号内的空白
    - 列表缩进/符号规范
    若三节都不存在，兜底输出一节“## 漏洞识别”+ 前文片段（不造内容）。
    """
    if not analysis_text:
        return ""

    s = html.unescape(analysis_text)
    s = _norm_newlines(s)
    s = _fix_heading_spaces(s)
    s = _trim_md_backticks(s)
    s = _ensure_code_fences(s)

    sections = []
    # 依次取“第一次出现”的三节内容，后续重复的同名节会被忽略（自然达到去重效果）
    for key in ["vuln", "impact", "reco"]:
        content = _find_first_section(s, TITLE_ALIASES[key])
        if content:
            content = _trim_md_backticks(content)
            content = _ensure_code_fences(content)
            content = _clean_list_indentation(content)
            # 标题行后插入一个空行，避免“标题与内容同一行”
            sections.append(f"{NORM_TITLES[key]}\n\n{content.strip()}\n")

    if not sections:
        # 兜底：取前 30 行，做最小可渲染片段
        snippet = "\n".join(s.split("\n")[:30]).strip()
        snippet = _clean_list_indentation(_ensure_code_fences(_trim_md_backticks(snippet)))
        return f"## 漏洞识别\n\n{snippet}\n"

    # 最终拼接，去除多余空白
    clean = "\n".join(sections).strip() + "\n"
    # 再次修补可能遗漏的标题换行（极端情况下）
    clean = _fix_heading_spaces(clean)
    return clean
